
from types import GeneratorType

from lingospot.misc import only
from lingospot.frozendict import frozendict

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------

reserved = [ 'n', 'partial' ]

def mapper( singular, plural = None ):
  if plural == None:
    plural = singular + 's'

  assert singular not in reserved
  assert plural not in reserved
  
  def wrapper( get_outputs ):
    get_outputs = decorate_get_outputs( plural, get_outputs )
    def map_either( *a, **kw ):
      n = get_n( kw )
      partial = get_partial( kw )        
      validate_args( a, kw )    
      input_arg = get_input_arg( singular, plural, n, a, kw )
      inputs = get_inputs( n, input_arg )
      inputs__outputs = call_get_outputs( plural, inputs, get_outputs, kw )
      validate_outputs( partial, inputs, inputs__outputs )
      return get_output_arg( n, partial, inputs, inputs__outputs ) 
    return map_either  
  return wrapper 

# --------------------------------------

class Invalid_input( Exception ): pass
class Invalid_output( Exception ): pass
class Partial( Invalid_output ): pass

mapper.Invalid_input = Invalid_input
mapper.Invalid_output = Invalid_output
mapper.Partial = Partial

# ------------------------------------------------------------------------------

def decorate_get_outputs( plural, get_outputs ):  
  def dict_wrapped( *a, **kw ):
    x = get_outputs( *a, **kw )
    if isinstance( x, GeneratorType ):
      return frozendict( x )
    else:
      return x
  
  return dict_wrapped

# ------------------------------------------------------------------------------
  
def get_n( kw ):        
  n = kw.pop( 'n', False )
  return n
 
# ------------------------------------------------------------------------------
  
def get_partial( kw ):        
  partial = kw.pop( 'partial', False )
  return partial

# ------------------------------------------------------------------------------

def validate_args( a, kw ):
      
  pos_count = len( a )
  kw_count = len( kw )
      
  if pos_count > 1:
    raise Invalid_input( 'too many positional inputs' )
  
  if pos_count == 0 and kw_count == 0:
    raise Invalid_input( 'insufficient inputs' )
    
# ------------------------------------------------------------------------------

def get_input_arg( singular, plural, n, a, kw ):
  pos_count = len( a )
  assert pos_count <= 1, 'impossible'
  
  if pos_count == 1:
    if singular in kw or plural in kw:
      m = 'cannot pass in keywords %s or %s with a positional arg' % ( singular, plural )
      raise Invalid_input( m )
    
    return only( a )
  
  elif n:
    if plural not in kw:
      raise Invalid_input( 'expecting keyword arg (%s)' % plural )
    
    if singular in kw:
      raise Invalid_input( 'not expecting keyword arg (%s)' % singular )
    
    return kw.pop( plural )
    
  else:
    if singular not in kw:
      raise Invalid_input( 'expecting keyword arg (%s)' % singular )
    
    if plural in kw:
      raise Invalid_input( 'not expecting keyword arg (%s)' % plural )
            
    return kw.pop( singular )

# ------------------------------------------------------------------------------

def get_inputs( n, input_arg ):
  if not n:
    return frozenset([ input_arg ])
  
  if isinstance( input_arg, set ):
    return input_arg
  
  return frozenset( input_arg )

# ------------------------------------------------------------------------------

def call_get_outputs( plural, inputs, get_outputs, kw ):

  assert plural not in kw, 'impossible'
  kw[ plural ] = inputs

  return get_outputs( **kw )
  
# ------------------------------------------------------------------------------

def validate_outputs( partial, inputs, inputs__outputs ):
  # validate that all inputs in inputs__outputs exist in inputs
  if set( inputs__outputs.keys() ) - set( inputs ):
    keys = set( inputs__outputs.keys() ) - set( inputs )
    m = 'got unexpected keys in mapper: %s' % ( keys, )
    raise Invalid_output( m )  # TODO: test this case
  
  if partial: return
  
  input_count = len( inputs )
  output_count = len( inputs__outputs )
  
  if input_count != output_count:
    m = 'input_count (%s) does not match output_count (%s)\ninputs = %s\noutputs = %s' % (
      input_count,
      output_count,
      inputs,
      inputs__outputs.values()
    )
    
    if input_count > output_count:
      raise Partial( m )
    else:
      raise Invalid_output( m )

# ------------------------------------------------------------------------------

def get_output_arg( n, partial, inputs, inputs__outputs ):
  if n:
    return inputs__outputs
  
  input = only( inputs )  
  
  if input in inputs__outputs.keys():
    return inputs__outputs[ input ]
  elif partial:
    return None
  else:
    m = 'no output for input (%s)' % input
    raise Invalid_output( m )
